"""
3.	LeNet-5由LeCun等人于1998年提出，主要进行手写数字识别，是经典的卷积神经网络。LeNet-5结构虽小，但卷积层、池化层、全连接层各模块齐全，
是学习 CNN的基础。使用深度学习框架tensorflow2.0+keras，完成下列各项要求内容（30分）：
"""
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import losses, optimizers, metrics, layers, activations
import os
import cv2 as cv
from sklearn.model_selection import train_test_split

# params
ALPHA = 0.001
N_CLS = 3
BATCH_SIZE = 64
N_EPOCHS = 10

np.random.seed(777)
tf.random.set_seed(777)

# ①	导入飞、汽车、鸟的数据集（data3）
# ②	正确读入图片的文件信息
x = []
y = []
dir = 'data/data3'
names = os.listdir(dir)
for name in names:
    path = os.path.join(dir, name)
    img = cv.imread(path, cv.IMREAD_COLOR)
    x.append(img)
    label = int(name[0])
    y.append(label)
x = np.float32(x) / 255.
y = np.float32(y)

# ③	获取文件标签，进行独热编码（这里3分类）
y = keras.utils.to_categorical(y, N_CLS)
print(f'x: {x.shape}, y: {y.shape}')

# ⑤	数据集切分为训练集和测试集
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=777)

# ④	进行洗牌乱序处理，自定义批量函数，实现每次处理一批数据
ds_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))\
    .shuffle(1000).batch(BATCH_SIZE).prefetch(tf.data.experimental.AUTOTUNE)
ds_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))\
    .shuffle(1000).batch(BATCH_SIZE).prefetch(tf.data.experimental.AUTOTUNE)

# ⑥	完成数据分析后，进入模型搭建工作（参照下图Lenet5）
# ⑦	要求模型由两层卷积，两层池化，最后3层全连接网络
inputs = keras.Input((32, 32, 3))
x = layers.Conv2D(6, (5, 5), (1, 1), 'valid')(inputs)
x = layers.MaxPooling2D((2, 2), (2, 2), 'same')(x)
x = layers.Conv2D(16, (5, 5), (1, 1), 'valid')(x)
x = layers.MaxPooling2D((2, 2), (2, 2), 'same')(x)
x = layers.Flatten()(x)
x = layers.Dense(400, activation=activations.relu)(x)
x = layers.Dense(120, activation=activations.relu)(x)
x = layers.Dense(84, activation=activations.relu)(x)
x = layers.Dense(N_CLS, activation=activations.softmax)(x)
model = keras.Model(inputs, x)
model.summary()
model.compile(
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    loss=losses.categorical_crossentropy,
    metrics=metrics.categorical_accuracy
)

# ⑧	训练集进行模型训练，要求每次大循环结束输出loss值
model.fit(ds_train, epochs=N_EPOCHS)

# ⑨	加入测试集数据，计算模型准确率
model.evaluate(ds_test)

# ⑩	随机抽取图片进行测试
N_SELECT = 10
print(f'随机抽取{N_SELECT}图片进行测试')
M_TEST = len(x_test)
rnd_idx = np.random.permutation(M_TEST)
select_idx = rnd_idx[:N_SELECT]  # ATTENTION [:10] vs [10]
x_selected = x_test[select_idx]
y_selected = y_test[select_idx]

pred = model.predict(x_selected)
y_label = np.argmax(y_selected, axis=1)
pred_label = np.argmax(pred, axis=1)
for pr, y_prob, pr_label, label in zip(pred, y_test, pred_label, y_label):
    print(f'{y_prob} => {pr}, {label} => {pr_label} ({label == pr_label})')
